# visual_verify_navi.py

import os
import os.path as osp
import sys
import glob
import logging
import argparse
import numpy as np
from tqdm import tqdm
import json
from scipy.spatial.transform import Rotation
import cv2
# bad
# --- Import from our visual utility file ---
from visual_util import launch_visualizer

# --- Import from the original utility file for data loading helpers ---
try:
    from verify_utils import setup_logging, depthmap_to_absolute_camera_coordinates, imread_cv2
except ImportError:
    print("FATAL ERROR: Could not import from 'verify_utils.py'.")
    print("Please ensure 'verify_utils.py' is in the same directory.")
    sys.exit(1)


def load_navi_view_data(rgb_path, annotations, logger, use_mask=True, downsample_factor=4):
    """
    Loads all necessary data for a single view from the NAVI dataset format.
    """
    try:
        filename = osp.basename(rgb_path)
        if filename not in annotations:
            logger.warning(f"No annotation found for {filename}. Skipping.")
            return None

        view_data = annotations[filename]
        camera_params = view_data['camera']
        
        base_dir = osp.dirname(osp.dirname(rgb_path))
        depth_path = osp.join(base_dir, "depth", filename.replace(".jpg", ".png"))
        mask_path = osp.join(base_dir, "masks", filename.replace(".jpg", ".png"))

        if not all(osp.exists(p) for p in [depth_path, mask_path]):
            logger.warning(f"Missing depth or mask file for {filename}. Skipping.")
            return None

        rgb_image = imread_cv2(rgb_path)
        depthmap = imread_cv2(depth_path, -1)
        
        if rgb_image is None or depthmap is None:
            logger.warning(f"Failed to read image or depth for {rgb_path}. Skipping.")
            return None
      
        rgb_image = rgb_image.astype(np.float32) / 255.0
        depthmap = depthmap.astype(np.float32) / 1000.0
        
        if use_mask:
            object_mask = imread_cv2(mask_path, 0)
            if object_mask is not None:
                object_mask = (object_mask.astype(np.float32) / 255.0) > 0.1
                depthmap *= object_mask
            else:
                logger.warning(f"Could not load mask file {mask_path}, proceeding without it.")
        
        if downsample_factor > 1:
            new_w = rgb_image.shape[1] // downsample_factor
            new_h = rgb_image.shape[0] // downsample_factor
            rgb_image = cv2.resize(rgb_image, (new_w, new_h), interpolation=cv2.INTER_AREA)
            depthmap = cv2.resize(depthmap, (new_w, new_h), interpolation=cv2.INTER_NEAREST)
            # if object_mask is not None:
            #     object_mask = cv2.resize(object_mask, (new_w, new_h), interpolation=cv2.INTER_NEAREST)

        focal_length = camera_params['focal_length']
        img_h, img_w = view_data['image_size'][1], view_data['image_size'][0]
        cx, cy = (img_w-1) / 2.0, (img_h-1) / 2.0
        intrinsics = np.array([
            [focal_length, 0, cx],
            [0, focal_length, cy],
            [0, 0, 1]
        ], dtype=np.float32)

        if downsample_factor > 1:
            intrinsics[:2, :] /= downsample_factor

        # --- Construct Camera Pose ---
        # The NAVI dataset provides the quaternion as [qw, qx, qy, qz].
        q_navi = np.array(camera_params['q'])
        
        # --- FIX: Reorder the quaternion for Scipy ---
        # Scipy's from_quat expects [x, y, z, w].
        q_scipy = np.array([q_navi[1], q_navi[2], q_navi[3], q_navi[0]])
        R = Rotation.from_quat(q_scipy).as_matrix()
        
        # Get translation vector and scale from mm to m
        t = np.array(camera_params['t']) / 1000.0
        
        # Assemble the camera-to-world (c2w) pose matrix
        camera_pose = np.eye(4, dtype=np.float32)
        camera_pose[:3, :3] = R
        camera_pose[:3, 3] = t
        # camera_pose = np.linalg.inv(camera_pose)

        return {
            "rgb": rgb_image, "depth": depthmap, "intrinsics": intrinsics,
            "pose": camera_pose, "path": rgb_path
        }
    except Exception as e:
        logger.error(f"Could not process view for {osp.basename(rgb_path)}. Details: {e}", exc_info=True)
        return None

# ======================================================================================
# --- Main Execution Block ---
# ======================================================================================

if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        description="Visually verify the consistency of a NAVI dataset."
    )
    
    parser.add_argument(
        "--data_root", 
        type=str,
        default='/lc/data/3D/navi/navi_v1.5',
        help="Path to the root directory of the NAVI dataset (e.g., /path/to/navi_v1.5)."
    )
    
    parser.add_argument(
        "--log_file", 
        type=str, 
        default="logs/visual_verification_navi.log",
        help="Path to the output log file."
    )

    parser.add_argument(
        "--scene", 
        type=str, 
        default='duck_bath_yellow_s',
        help="Optional: Specify a single object ID (e.g., 'duck_bath_yellow_s') to visualize."
    )
    parser.add_argument(
        "--subscene", 
        type=str, 
        default='multiview-02-pixel_6pro',
        help="Optional: Specify a single object ID (e.g., 'duck_bath_yellow_s') to visualize."
    )

    parser.add_argument(
        "--port", 
        type=int, 
        default=8080,
        help="Port to run the Viser web server on."
    )
    
    parser.add_argument(
        "--no_mask",
        action="store_true",
        help="If set, do not apply the foreground mask to the depth map."
    )
    
    args = parser.parse_args()
    
    log_dir = osp.dirname(args.log_file)
    if log_dir:
        os.makedirs(log_dir, exist_ok=True)
    logger = setup_logging(args.log_file)

    if not osp.isdir(args.data_root):
        logger.error(f"Provided data root '{args.data_root}' does not exist.")
        sys.exit(1)
    
    if args.scene:
        top_level_scenes = [osp.join(args.data_root, args.scene)]
        if not osp.isdir(top_level_scenes[0]):
            logger.error(f"Specified scene directory '{top_level_scenes[0]}' does not exist.")
            sys.exit(1)
    else:
        top_level_scenes = sorted([d for d in glob.glob(osp.join(args.data_root, '*')) if osp.isdir(d)])
    
    logger.info(f"Found {len(top_level_scenes)} top-level scenes to process in '{args.data_root}'.")
    
    for scene_path in top_level_scenes:
        subscene_paths = sorted([d for d in glob.glob(osp.join(scene_path, 'multiview-*')) if osp.isdir(d)])
        for subscene_path in subscene_paths:
            scene_key = f"{osp.basename(scene_path)}/{osp.basename(subscene_path)}"
            logger.info(f"\n--- Processing subscene: {scene_key} ---")

            # --- Load annotations for the entire subscene once ---
            annotations_path = osp.join(subscene_path, "annotations.json")
            if not osp.exists(annotations_path):
                logger.warning(f"No annotations.json found in {subscene_path}. Skipping.")
                continue
            
            with open(annotations_path, 'r') as f:
                annotations_list = json.load(f)
            
            # Create a lookup dictionary for efficient access
            annotations_data = {item['filename']: item for item in annotations_list}

            image_dir = osp.join(subscene_path, "images")
            rgb_paths = sorted(glob.glob(osp.join(image_dir, "*.jpg")))
            if not rgb_paths:
                logger.warning(f"No JPG images found in {image_dir}. Skipping.")
                continue

            pc_list, color_list, conf_list = [], [], []
            focals, pps, Rs, ts = [], [], [], []
            
            logger.info("Loading views and generating point clouds...")
            for rgb_path in tqdm(rgb_paths, file=sys.stdout):
                view_data = load_navi_view_data(rgb_path, annotations_data, logger, use_mask=not args.no_mask)
                if view_data:
                    world_points, valid_mask, _ = depthmap_to_absolute_camera_coordinates(
                        view_data["depth"], view_data["intrinsics"], view_data["pose"]
                    )
                    pc_list.append(world_points)
                    color_list.append(view_data['rgb'])
                    conf_list.append(valid_mask.astype(np.float32))

                    intrinsics = view_data['intrinsics']
                    pose = view_data['pose']
                    focals.append(intrinsics[0, 0])
                    pps.append((intrinsics[0, 2], intrinsics[1, 2]))
                    Rs.append(pose[:3, :3])
                    ts.append(pose[:3, 3])
        
            if not pc_list:
                logger.warning(f"No valid data could be loaded for {scene_key}. Skipping.")
                continue

            cam_dict = {'focal': focals, 'pp': pps, 'R': Rs, 't': ts}

            try:
                logger.info(f"Launching interactive visualizer for {scene_key} at http://localhost:{args.port}")
                logger.info("--> Close the browser tab or press Ctrl+C in the console to continue.")
                launch_visualizer(
                    pc_list=pc_list, color_list=color_list, conf_list=conf_list,
                    cam_dict=cam_dict, port=args.port
                )
            except KeyboardInterrupt:
                logger.info(f"\nVisualizer for {scene_key} closed, proceeding...")
            except Exception as e:
                logger.error(f"Failed to launch visualizer for {scene_key}: {e}", exc_info=True)

    logger.info("\n--- Visual verification complete. ---")
